import websocket
import thread
import time


class WebSocketServer:
    def __init__(self, url):
        self.url = url  # WebSocket服务器的URL
        self.ws = None  # WebSocket连接实例
        self.keep_running = True  # 控制运行状态的标志

    def on_message(self, message):
        # 当从客户端接收到消息时执行此函数
        print("Received frame")
        # 此处可以添加代码将接收到的消息（视频帧数据）发送给MediaPipe处理模块

    def on_error(self, error):
        # 当WebSocket出现错误时执行此函数
        print("Error:", error)

    def on_close(self):
        # 当WebSocket连接关闭时执行此函数
        print("### closed ###")

    def on_open(self):
        # 当WebSocket连接成功开启时执行此函数
        def run(*args):
            while self.keep_running:
                time.sleep(1)  # 模拟持续运行
            self.ws.close()  # 关闭WebSocket连接
            print("Thread terminating...")

        thread.start_new_thread(run, ())  # 开启新线程运行

    def start(self):
        # 启动WebSocket服务器
        self.ws = websocket.WebSocketApp(self.url,
                                         on_message=self.on_message,
                                         on_error=self.on_error,
                                         on_close=self.on_close)
        self.ws.on_open = self.on_open
        self.ws.run_forever()  # 持续运行WebSocket连接

    def stop(self):
        # 停止WebSocket服务器
        self.keep_running = False  # 更改运行状态标志
